"""
    ModelInterface

Multi-model ensemble I/O interface for the GSI/EnKF system, based on gridio_*.f90 modules.
Provides unified interface for reading/writing ensemble member files across different models
(GFS, WRF, FV3, NMMB) with proper metadata handling and format conversion.

Key Features:
- Multi-model ensemble I/O (GFS, WRF, FV3, NMMB)
- Ensemble member reading/writing with metadata preservation
- State vector initialization and finalization
- Model-specific grid configurations and transformations
- Ensemble forecast integration and management
- Format conversion and data validation

Based on:
- gridio_gfs.f90 - GFS model interface
- gridio_wrf.f90 - WRF model interface  
- gridio_nmmb.f90 - NMMB model interface
- gridinfo_*.f90 - Grid information modules
"""
module ModelInterface

using LinearAlgebra
using SparseArrays
using Printf
using Dates
using NetCDF
using HDF5

export AbstractModelInterface, ModelConfig
export GFSInterface, WRFInterface, FV3Interface, NMMBInterface
export read_ensemble_member, write_ensemble_member
export initialize_model_interface, finalize_model_interface
export convert_state_vector, validate_ensemble_data
export get_model_grid_info, transform_coordinates

# Abstract interfaces
abstract type AbstractModelInterface end

"""
    ModelConfig

Configuration structure for model-specific parameters and settings.
"""
struct ModelConfig
    model_type::String                    # "GFS", "WRF", "FV3", "NMMB"
    grid_type::String                    # "gaussian", "latlon", "cubed_sphere", "rotated"
    nlons::Int                           # Number of longitude points
    nlats::Int                           # Number of latitude points  
    nlevs::Int                           # Number of vertical levels
    data_path::String                    # Path to model data files
    ensemble_size::Int                   # Number of ensemble members
    time_levels::Int                     # Number of time levels
    reduced_grid::Bool                   # Use reduced gaussian grid
    hybrid_vertical::Bool                # Use hybrid coordinates
    spectral_transform::Bool             # Use spectral transform
    file_format::String                  # "nemsio", "netcdf", "binary", "grib"
    variable_names::Vector{String}       # Model variable names
    coordinate_transform::Dict{String,Any} # Coordinate transformation parameters
end

"""
    GFSInterface

Interface for NCEP GFS model ensemble I/O operations.
Based on gridio_gfs.f90 functionality.
"""
struct GFSInterface <: AbstractModelInterface
    config::ModelConfig
    spectral_coeffs::Dict{String,Array{Float64}}
    pressure_levels::Vector{Float64}
    sigma_levels::Vector{Float64}
    surface_pressure::Array{Float64}
    
    function GFSInterface(config::ModelConfig)
        @assert config.model_type == "GFS"
        spectral_coeffs = Dict{String,Array{Float64}}()
        pressure_levels = zeros(config.nlevs)
        sigma_levels = zeros(config.nlevs + 1)
        surface_pressure = zeros(config.nlons, config.nlats)
        new(config, spectral_coeffs, pressure_levels, sigma_levels, surface_pressure)
    end
end

"""
    WRFInterface

Interface for WRF model ensemble I/O operations.
Based on gridio_wrf.f90 functionality.
"""
struct WRFInterface <: AbstractModelInterface
    config::ModelConfig
    eta_levels::Vector{Float64}
    map_projection::Dict{String,Float64}
    domain_boundaries::Dict{String,Float64}
    mass_points::Array{Float64,2}
    
    function WRFInterface(config::ModelConfig)
        @assert config.model_type == "WRF"
        eta_levels = zeros(config.nlevs + 1)
        map_projection = Dict("lat_1" => 30.0, "lat_2" => 60.0, "lat_0" => 40.0, "lon_0" => -97.0)
        domain_boundaries = Dict("west" => -130.0, "east" => -60.0, "south" => 20.0, "north" => 55.0)
        mass_points = zeros(config.nlons, config.nlats)
        new(config, eta_levels, map_projection, domain_boundaries, mass_points)
    end
end

"""
    FV3Interface

Interface for FV3 (Finite Volume Cubed-Sphere) model ensemble I/O.
"""
struct FV3Interface <: AbstractModelInterface
    config::ModelConfig
    cubed_sphere_grid::Dict{String,Array{Float64}}
    ak_bk_coefficients::Dict{String,Vector{Float64}}
    tile_layout::Dict{String,Int}
    
    function FV3Interface(config::ModelConfig)
        @assert config.model_type == "FV3"
        cubed_sphere_grid = Dict{String,Array{Float64}}()
        ak_bk_coefficients = Dict("ak" => zeros(config.nlevs+1), "bk" => zeros(config.nlevs+1))
        tile_layout = Dict("ntiles" => 6, "nx_tile" => config.nlons÷6, "ny_tile" => config.nlats÷6)
        new(config, cubed_sphere_grid, ak_bk_coefficients, tile_layout)
    end
end

"""
    NMMBInterface

Interface for NMMB (Nonhydrostatic Multiscale Model on the B-grid) ensemble I/O.
Based on gridio_nmmb.f90 functionality.
"""
struct NMMBInterface <: AbstractModelInterface
    config::ModelConfig
    b_grid_staggering::Dict{String,Array{Float64}}
    vertical_coordinates::Dict{String,Vector{Float64}}
    mass_flux_variables::Dict{String,Array{Float64}}
    
    function NMMBInterface(config::ModelConfig)
        @assert config.model_type == "NMMB"
        b_grid_staggering = Dict{String,Array{Float64}}()
        vertical_coordinates = Dict("eta" => zeros(config.nlevs), "zeta" => zeros(config.nlevs))
        mass_flux_variables = Dict{String,Array{Float64}}()
        new(config, b_grid_staggering, vertical_coordinates, mass_flux_variables)
    end
end

"""
    initialize_model_interface(model_type::String, config_params::Dict)

Initialize model interface based on model type and configuration parameters.
"""
function initialize_model_interface(model_type::String, config_params::Dict)
    # Create model configuration
    config = ModelConfig(
        model_type,
        get(config_params, "grid_type", "gaussian"),
        get(config_params, "nlons", 360),
        get(config_params, "nlats", 180),
        get(config_params, "nlevs", 64),
        get(config_params, "data_path", "./"),
        get(config_params, "ensemble_size", 40),
        get(config_params, "time_levels", 1),
        get(config_params, "reduced_grid", false),
        get(config_params, "hybrid_vertical", true),
        get(config_params, "spectral_transform", true),
        get(config_params, "file_format", "nemsio"),
        get(config_params, "variable_names", ["u", "v", "t", "q", "ps"]),
        get(config_params, "coordinate_transform", Dict{String,Any}())
    )
    
    # Create appropriate interface
    if model_type == "GFS"
        interface = GFSInterface(config)
        _initialize_gfs_specifics!(interface)
    elseif model_type == "WRF"
        interface = WRFInterface(config)
        _initialize_wrf_specifics!(interface)
    elseif model_type == "FV3"
        interface = FV3Interface(config)
        _initialize_fv3_specifics!(interface)
    elseif model_type == "NMMB"
        interface = NMMBInterface(config)
        _initialize_nmmb_specifics!(interface)
    else
        error("Unsupported model type: $model_type")
    end
    
    println("Initialized $(model_type) model interface")
    return interface
end

"""
    read_ensemble_member(interface::AbstractModelInterface, member_id::Int, time_level::Int=1)

Read ensemble member data from file with proper format handling and metadata extraction.
"""
function read_ensemble_member(interface::AbstractModelInterface, member_id::Int, time_level::Int=1)
    config = interface.config
    filename = _construct_filename(config, member_id, time_level)
    
    # Check if file exists
    if !isfile(filename)
        error("Ensemble member file not found: $filename")
    end
    
    println("Reading ensemble member $member_id from $filename")
    
    # Read data based on file format
    if config.file_format == "nemsio"
        return _read_nemsio_data(interface, filename)
    elseif config.file_format == "netcdf"
        return _read_netcdf_data(interface, filename)
    elseif config.file_format == "binary"
        return _read_binary_data(interface, filename)
    elseif config.file_format == "grib"
        return _read_grib_data(interface, filename)
    else
        error("Unsupported file format: $(config.file_format)")
    end
end

"""
    write_ensemble_member(interface::AbstractModelInterface, member_id::Int, data::Dict, time_level::Int=1)

Write ensemble member data to file with proper format and metadata handling.
"""
function write_ensemble_member(interface::AbstractModelInterface, member_id::Int, data::Dict, time_level::Int=1)
    config = interface.config
    filename = _construct_filename(config, member_id, time_level, write=true)
    
    println("Writing ensemble member $member_id to $filename")
    
    # Validate data before writing
    _validate_ensemble_data(interface, data)
    
    # Write data based on file format
    if config.file_format == "nemsio"
        _write_nemsio_data(interface, filename, data)
    elseif config.file_format == "netcdf"
        _write_netcdf_data(interface, filename, data)
    elseif config.file_format == "binary"
        _write_binary_data(interface, filename, data)
    elseif config.file_format == "grib"
        _write_grib_data(interface, filename, data)
    else
        error("Unsupported file format: $(config.file_format)")
    end
    
    println("Successfully wrote ensemble member $member_id")
    return true
end

"""
    convert_state_vector(interface::AbstractModelInterface, state_vector::Vector{Float64}, 
                         grid_data::Dict)

Convert between state vector representation and grid data format.
"""
function convert_state_vector(interface::AbstractModelInterface, state_vector::Vector{Float64}, 
                             grid_data::Dict)
    config = interface.config
    
    if typeof(interface) == GFSInterface
        return _convert_gfs_state_vector(interface, state_vector, grid_data)
    elseif typeof(interface) == WRFInterface
        return _convert_wrf_state_vector(interface, state_vector, grid_data)
    elseif typeof(interface) == FV3Interface
        return _convert_fv3_state_vector(interface, state_vector, grid_data)
    elseif typeof(interface) == NMMBInterface
        return _convert_nmmb_state_vector(interface, state_vector, grid_data)
    else
        error("Unsupported interface type")
    end
end

"""
    validate_ensemble_data(interface::AbstractModelInterface, data::Dict)

Validate ensemble data for consistency and physical constraints.
"""
function validate_ensemble_data(interface::AbstractModelInterface, data::Dict)
    config = interface.config
    
    # Basic validation checks
    validation_results = Dict{String,Bool}()
    
    # Check required variables
    required_vars = config.variable_names
    for var in required_vars
        if haskey(data, var)
            validation_results[var * "_present"] = true
        else
            validation_results[var * "_present"] = false
            @warn "Missing required variable: $var"
        end
    end
    
    # Check data dimensions
    for (var_name, var_data) in data
        if isa(var_data, Array)
            expected_dims = _get_expected_dimensions(interface, var_name)
            if size(var_data) != expected_dims
                validation_results[var_name * "_dimensions"] = false
                @warn "Incorrect dimensions for $var_name: $(size(var_data)) vs $expected_dims"
            else
                validation_results[var_name * "_dimensions"] = true
            end
        end
    end
    
    # Model-specific validation
    if typeof(interface) == GFSInterface
        _validate_gfs_data(interface, data, validation_results)
    elseif typeof(interface) == WRFInterface
        _validate_wrf_data(interface, data, validation_results)
    elseif typeof(interface) == FV3Interface
        _validate_fv3_data(interface, data, validation_results)
    elseif typeof(interface) == NMMBInterface
        _validate_nmmb_data(interface, data, validation_results)
    end
    
    # Check if all validations passed
    all_passed = all(values(validation_results))
    if !all_passed
        failed_checks = [k for (k,v) in validation_results if !v]
        error("Data validation failed for: $(join(failed_checks, ", "))")
    end
    
    println("Data validation passed for all checks")
    return validation_results
end

"""
    get_model_grid_info(interface::AbstractModelInterface)

Get comprehensive grid information for the model.
"""
function get_model_grid_info(interface::AbstractModelInterface)
    config = interface.config
    
    grid_info = Dict{String,Any}(
        "model_type" => config.model_type,
        "grid_type" => config.grid_type,
        "nlons" => config.nlons,
        "nlats" => config.nlats,
        "nlevs" => config.nlevs,
        "reduced_grid" => config.reduced_grid,
        "hybrid_vertical" => config.hybrid_vertical
    )
    
    # Add model-specific grid information
    if typeof(interface) == GFSInterface
        grid_info["spectral_truncation"] = get(config.coordinate_transform, "spectral_truncation", 574)
        grid_info["sigma_levels"] = interface.sigma_levels
    elseif typeof(interface) == WRFInterface
        grid_info["map_projection"] = interface.map_projection
        grid_info["domain_boundaries"] = interface.domain_boundaries
        grid_info["eta_levels"] = interface.eta_levels
    elseif typeof(interface) == FV3Interface
        grid_info["tile_layout"] = interface.tile_layout
        grid_info["ak_bk_coefficients"] = interface.ak_bk_coefficients
    elseif typeof(interface) == NMMBInterface
        grid_info["b_grid_staggering"] = "B-grid staggered"
        grid_info["vertical_coordinates"] = interface.vertical_coordinates
    end
    
    return grid_info
end

"""
    transform_coordinates(interface::AbstractModelInterface, data::Dict, source_coords::String, 
                         target_coords::String)

Transform data between different coordinate systems.
"""
function transform_coordinates(interface::AbstractModelInterface, data::Dict, 
                              source_coords::String, target_coords::String)
    if source_coords == target_coords
        return data  # No transformation needed
    end
    
    transformed_data = deepcopy(data)
    
    if typeof(interface) == GFSInterface
        return _transform_gfs_coordinates(interface, transformed_data, source_coords, target_coords)
    elseif typeof(interface) == WRFInterface
        return _transform_wrf_coordinates(interface, transformed_data, source_coords, target_coords)
    elseif typeof(interface) == FV3Interface
        return _transform_fv3_coordinates(interface, transformed_data, source_coords, target_coords)
    elseif typeof(interface) == NMMBInterface
        return _transform_nmmb_coordinates(interface, transformed_data, source_coords, target_coords)
    else
        error("Coordinate transformation not implemented for this interface type")
    end
end

"""
    finalize_model_interface(interface::AbstractModelInterface)

Clean up and finalize model interface resources.
"""
function finalize_model_interface(interface::AbstractModelInterface)
    println("Finalizing $(interface.config.model_type) model interface")
    
    # Model-specific cleanup
    if typeof(interface) == GFSInterface
        _finalize_gfs_interface(interface)
    elseif typeof(interface) == WRFInterface
        _finalize_wrf_interface(interface)
    elseif typeof(interface) == FV3Interface
        _finalize_fv3_interface(interface)
    elseif typeof(interface) == NMMBInterface
        _finalize_nmmb_interface(interface)
    end
    
    println("Model interface finalized successfully")
    return true
end

# Private helper functions
function _construct_filename(config::ModelConfig, member_id::Int, time_level::Int; write::Bool=false)
    member_str = @sprintf("mem%03d", member_id)
    time_str = @sprintf("t%02d", time_level)
    
    if config.model_type == "GFS"
        ext = config.file_format == "nemsio" ? ".nemsio" : ".nc"
        return joinpath(config.data_path, "$(member_str)_gfs_$(time_str)$(ext)")
    elseif config.model_type == "WRF"
        return joinpath(config.data_path, "wrfout_d01_$(member_str)_$(time_str).nc")
    elseif config.model_type == "FV3"
        return joinpath(config.data_path, "fv3_$(member_str)_$(time_str).nc")
    elseif config.model_type == "NMMB"
        return joinpath(config.data_path, "nmmb_$(member_str)_$(time_str).nc")
    else
        error("Unknown model type for filename construction")
    end
end

function _get_expected_dimensions(interface::AbstractModelInterface, var_name::String)
    config = interface.config
    
    # 3D variables (u, v, t, q, etc.)
    if var_name in ["u", "v", "t", "q", "oz", "cw"]
        return (config.nlons, config.nlats, config.nlevs)
    # 2D variables (ps, sst, etc.)
    elseif var_name in ["ps", "sst", "pst"]
        return (config.nlons, config.nlats)
    # Pressure levels
    elseif var_name == "prse"
        return (config.nlons, config.nlats, config.nlevs + 1)
    else
        # Default to 3D
        return (config.nlons, config.nlats, config.nlevs)
    end
end

# Model-specific initialization functions (simplified implementations)
function _initialize_gfs_specifics!(interface::GFSInterface)
    # Initialize GFS-specific parameters
    config = interface.config
    interface.sigma_levels[:] = range(0.0, 1.0, length=config.nlevs+1)
    interface.pressure_levels[:] = range(1000.0, 1.0, length=config.nlevs)
    println("GFS interface initialized with spectral transform capability")
end

function _initialize_wrf_specifics!(interface::WRFInterface)
    # Initialize WRF-specific parameters
    config = interface.config
    interface.eta_levels[:] = range(1.0, 0.0, length=config.nlevs+1)
    println("WRF interface initialized with eta coordinates")
end

function _initialize_fv3_specifics!(interface::FV3Interface)
    # Initialize FV3-specific parameters
    config = interface.config
    # Set up hybrid coordinates
    interface.ak_bk_coefficients["ak"][:] = range(0.0, 100.0, length=config.nlevs+1)
    interface.ak_bk_coefficients["bk"][:] = range(1.0, 0.0, length=config.nlevs+1)
    println("FV3 interface initialized with cubed-sphere grid")
end

function _initialize_nmmb_specifics!(interface::NMMBInterface)
    # Initialize NMMB-specific parameters  
    config = interface.config
    interface.vertical_coordinates["eta"][:] = range(1.0, 0.0, length=config.nlevs)
    interface.vertical_coordinates["zeta"][:] = range(0.0, 1.0, length=config.nlevs)
    println("NMMB interface initialized with B-grid staggering")
end

# Placeholder implementations for format-specific I/O (would need full implementation)
function _read_nemsio_data(interface::AbstractModelInterface, filename::String)
    # This would use NEMSIO library bindings to read data
    # For now, return dummy data structure
    config = interface.config
    return Dict{String,Any}(
        "u" => zeros(Float64, config.nlons, config.nlats, config.nlevs),
        "v" => zeros(Float64, config.nlons, config.nlats, config.nlevs),
        "t" => zeros(Float64, config.nlons, config.nlats, config.nlevs),
        "q" => zeros(Float64, config.nlons, config.nlats, config.nlevs),
        "ps" => zeros(Float64, config.nlons, config.nlats)
    )
end

function _read_netcdf_data(interface::AbstractModelInterface, filename::String)
    # Read NetCDF data
    config = interface.config
    data = Dict{String,Any}()
    
    ncopen(filename) do nc
        for var_name in config.variable_names
            if haskey(nc.vars, var_name)
                data[var_name] = nc[var_name][:,:]
            else
                @warn "Variable $var_name not found in NetCDF file"
            end
        end
    end
    return data
end

function _read_binary_data(interface::AbstractModelInterface, filename::String)
    # Binary data reading implementation
    error("Binary format reading not yet implemented")
end

function _read_grib_data(interface::AbstractModelInterface, filename::String)
    # GRIB data reading implementation  
    error("GRIB format reading not yet implemented")
end

function _write_nemsio_data(interface::AbstractModelInterface, filename::String, data::Dict)
    # NEMSIO writing implementation
    println("Would write NEMSIO data to $filename")
end

function _write_netcdf_data(interface::AbstractModelInterface, filename::String, data::Dict)
    # NetCDF writing implementation
    config = interface.config
    
    ncopen(filename, "w") do nc
        # Define dimensions
        defDim(nc, "lon", config.nlons)
        defDim(nc, "lat", config.nlats)
        defDim(nc, "lev", config.nlevs)
        
        # Define and write variables
        for (var_name, var_data) in data
            if ndims(var_data) == 3
                defVar(nc, var_name, Float64, ("lon", "lat", "lev"))
            elseif ndims(var_data) == 2
                defVar(nc, var_name, Float64, ("lon", "lat"))
            end
            nc[var_name][:] = var_data
        end
    end
end

function _write_binary_data(interface::AbstractModelInterface, filename::String, data::Dict)
    # Binary writing implementation
    error("Binary format writing not yet implemented")
end

function _write_grib_data(interface::AbstractModelInterface, filename::String, data::Dict)
    # GRIB writing implementation
    error("GRIB format writing not yet implemented")
end

# Model-specific validation functions (simplified)
function _validate_gfs_data(interface::GFSInterface, data::Dict, validation_results::Dict)
    # GFS-specific validation
    if haskey(data, "ps")
        ps = data["ps"]
        if any(ps .< 0) || any(ps .> 110000)
            validation_results["ps_physical_range"] = false
            @warn "Surface pressure out of physical range"
        else
            validation_results["ps_physical_range"] = true
        end
    end
end

function _validate_wrf_data(interface::WRFInterface, data::Dict, validation_results::Dict)
    # WRF-specific validation
    validation_results["wrf_specific"] = true
end

function _validate_fv3_data(interface::FV3Interface, data::Dict, validation_results::Dict)
    # FV3-specific validation
    validation_results["fv3_specific"] = true
end

function _validate_nmmb_data(interface::NMMBInterface, data::Dict, validation_results::Dict)
    # NMMB-specific validation
    validation_results["nmmb_specific"] = true
end

# Model-specific data conversion functions (simplified)
function _convert_gfs_state_vector(interface::GFSInterface, state_vector::Vector{Float64}, grid_data::Dict)
    # GFS state vector conversion with spectral transforms
    return grid_data  # Simplified
end

function _convert_wrf_state_vector(interface::WRFInterface, state_vector::Vector{Float64}, grid_data::Dict)
    # WRF state vector conversion
    return grid_data  # Simplified
end

function _convert_fv3_state_vector(interface::FV3Interface, state_vector::Vector{Float64}, grid_data::Dict)
    # FV3 cubed-sphere state vector conversion
    return grid_data  # Simplified  
end

function _convert_nmmb_state_vector(interface::NMMBInterface, state_vector::Vector{Float64}, grid_data::Dict)
    # NMMB B-grid state vector conversion
    return grid_data  # Simplified
end

# Coordinate transformation functions (simplified)
function _transform_gfs_coordinates(interface::GFSInterface, data::Dict, source::String, target::String)
    return data  # Simplified
end

function _transform_wrf_coordinates(interface::WRFInterface, data::Dict, source::String, target::String)
    return data  # Simplified
end

function _transform_fv3_coordinates(interface::FV3Interface, data::Dict, source::String, target::String)
    return data  # Simplified
end

function _transform_nmmb_coordinates(interface::NMMBInterface, data::Dict, source::String, target::String)
    return data  # Simplified
end

# Cleanup functions
function _finalize_gfs_interface(interface::GFSInterface)
    # GFS-specific cleanup
end

function _finalize_wrf_interface(interface::WRFInterface)
    # WRF-specific cleanup
end

function _finalize_fv3_interface(interface::FV3Interface)
    # FV3-specific cleanup
end

function _finalize_nmmb_interface(interface::NMMBInterface)
    # NMMB-specific cleanup
end

function _validate_ensemble_data(interface::AbstractModelInterface, data::Dict)
    # Internal validation function
    validate_ensemble_data(interface, data)
end

end # module ModelInterface